package org.zero;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.Modifier;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * 常用的字节码操作类库
 * <ul>
 *     <li>Javassist</li>
 *     <li>Byte-Buddy</li>
 *     <li>ASM</li>
 *     <li>CGLIB (Code Generation Library)</li>
 *     <li>BCEL (Apache Commons Bytecode Engineering Library)</li>
 *     <li>Objenesis</li>
 * </ul>
 *
 * @author zero
 * @since 2021/1/11
 */
public class CustomClassFileTransformer implements ClassFileTransformer {
    private static final String TEST_CLASS_NAME = "org.zero.Test";

    @Override
    public byte[] transform(ClassLoader loader,
                            String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {

        try {
            String replacedClassName = className.replace('/', '.');
            if (replacedClassName.equals(TEST_CLASS_NAME)) {
                final ClassPool classPool = ClassPool.getDefault();
                final CtClass ctClass = classPool.get(TEST_CLASS_NAME);

                for (CtMethod ctMethod : ctClass.getMethods()) {
                    // 跳过 Native 方法
                    if (Modifier.isNative(ctMethod.getModifiers())) {
                        continue;
                    }

                    String methodName = ctMethod.getName();

                    ctMethod.insertBefore("System.out.println(\"" + ctClass.getSimpleName() + "."
                            + methodName + " start.\");");
                    ctMethod.insertAfter("System.out.println(\"" + ctClass.getSimpleName() + "."
                            + methodName + " end.\");", false);
                }
                return ctClass.toBytecode();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }
}